import torch
from torch import nn
from torchsummary import summary

class Inception(nn.Module):
    def __init__(self, in_channels,c1,c2,c3,c4):#每个Inception通道数不一样,c1是卷积核的数量
        super(Inception, self).__init__()
        self.ReLU=nn.ReLU()

        #路线1，单1*1卷积层
        self.p1_1=nn.Conv2d(in_channels=in_channels, out_channels=c1, kernel_size=1)#步幅为1填充为0是默认

        #路线2,1*1卷积层，3*3的卷积
        self.p2_1=nn.Conv2d(in_channels=in_channels, out_channels=c2[0], kernel_size=1)
        self.p2_2 = nn.Conv2d(in_channels=c2[0], out_channels=c2[1], kernel_size=3,padding=1)

        #路径3,1*1卷积层，5*5的卷积。卷积stride默认为1
        self.p3_1 = nn.Conv2d(in_channels=in_channels, out_channels=c3[0], kernel_size=1)
        self.p3_2 = nn.Conv2d(in_channels=c3[0], out_channels=c3[1], kernel_size=5, padding=2)

        #路径4,3*3的最大池化，1*1的卷积
        #ALexNet用的是重叠池化，即若果kernel_size是3，stride是1
        self.p4_1=nn.MaxPool2d(kernel_size=3,padding=1, stride=1)#在池化中stride默认和kernel_size保持一致.卷积时stride默认为1。
        self.p4_2 = nn.Conv2d(in_channels=in_channels, out_channels=c4, kernel_size=1)

    def forward(self, x):
        p1 = self.ReLU(self.p1_1(x))
        p2 = self.ReLU(self.p2_2(self.ReLU(self.p2_1(x))))
        p3 = self.ReLU(self.p3_2(self.ReLU(self.p3_1(x))))
        p4 = self.ReLU(self.p4_2(self.p4_1(x)))
        return torch.cat((p1,p2,p3,p4),dim=1)#在卷积神经网络中，一个特征图的形状通常是 (batch_size, channels, height, width)


class GoogLeNet(nn.Module):#注意：由于LRN作用不明显，该模型没加LRN
    def __init__(self,Inception):
        super(GoogLeNet, self).__init__()
        #定义块。
        self.b1=nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=64, kernel_size=7,stride=2, padding=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )

        self.b2=nn.Sequential(
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=1, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(in_channels=64, out_channels=192, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )

        self.b3=nn.Sequential(
            Inception(in_channels=192, c1=64, c2=(96,128), c3=(16,32),c4=32),#c2是元组
            Inception(in_channels=256, c1=128, c2=(128,192), c3=(32,96),c4=64),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )

        self.b4=nn.Sequential(
            Inception(in_channels=480, c1=192, c2=(96, 208), c3=(16, 48), c4=64),  # c2是元组
            Inception(in_channels=512, c1=160, c2=(112, 224), c3=(24, 64), c4=64),
            Inception(in_channels=512, c1=128, c2=(128,256), c3=(24,64), c4=64),  # c2是元组
            Inception(in_channels=512, c1=112, c2=(128,288), c3=(32, 64), c4=64),
            Inception(in_channels=528, c1=256, c2=(160,320), c3=(32, 128), c4=128),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        )

        self.b5 = nn.Sequential(
            Inception(in_channels=832, c1=256, c2=(160,320), c3=(32,128), c4=128),  # c2是元组
            Inception(in_channels=832, c1=384, c2=(192,384), c3=(48,128), c4=128),
            nn.AdaptiveAvgPool2d((1,1)),#全局平均池化
            nn.Flatten(),
            nn.Linear(in_features=1024, out_features=10),
        )

        for m in self.modules():#赋初值
            if isinstance(m, nn.Conv2d):#如果是卷积
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)#偏置初始为0
                elif isinstance(m,nn.linear):#如果是全连接层
                    nn.init.normal_(m.weight, std=0.01)
                    if m.bias is not None:
                        nn.init.constant_(m.bias, 0)

    def forward(self,x):
        x=self.b1(x)
        x=self.b2(x)
        x=self.b3(x)
        x=self.b4(x)
        x=self.b5(x)

        return x

if __name__=='__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model=GoogLeNet(Inception).to(device)#前面讲的ALexNet和LeNet都没有参数
    print(summary(model,(1,224,224)))










